# -*- coding: utf-8 -*-
"""
Created on Mon Sep  6 19:10:49 2022
@author: hegang
"""
import os
import random
from PIL import Image
from torch.utils.data import  Dataset
from torchvision import transforms
import torch


class TripletData(Dataset):
    def __init__(self, path, transforms, split="train"):
        self.path = path
        self.split = split  # train or valid
        self.cats = 5  # number of categories
        self.transforms = transforms

    def __getitem__(self, idx):
        # our positive class for the triplet
        idx = str(idx % self.cats + 1)

        # choosing our pair of positive images (im1, im2)
        positives = os.listdir(os.path.join(self.path, idx))
        im1, im2 = random.sample(positives, 2)

        # choosing a negative class and negative image (im3)
        negative_cats = [str(x + 1) for x in range(self.cats)]
        negative_cats.remove(idx)
        negative_cat = str(random.choice(negative_cats))
        negatives = os.listdir(os.path.join(self.path, negative_cat))
        im3 = random.choice(negatives)

        im1, im2, im3 = os.path.join(self.path, idx, im1), os.path.join(self.path, idx, im2), os.path.join(self.path,
                                                                                                           negative_cat,
                                                                                                           im3)

        im1 = self.transforms(Image.open(im1))
        im2 = self.transforms(Image.open(im2))
        im3 = self.transforms(Image.open(im3))

        return [im1, im2, im3]

    # we'll put some value that we want since there can be far too many triplets possible
    # multiples of the number of images/ number of categories is a good choice
    def __len__(self):
        return self.cats * 8



if __name__ == '__main__':
    import matplotlib.pyplot as plt
    # Transforms
    train_transforms = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])

    PATH_TRAIN='/home/hegang/datas2/hegang/datas/public_datasets/contrast_learning/trainData'
    # Datasets and Dataloaders
    train_data = TripletData(PATH_TRAIN, train_transforms)
    train_loader = torch.utils.data.DataLoader(dataset=train_data, batch_size=1, shuffle=True, num_workers=4)

    for data in train_loader:
        im1,im2,im3=data
        plt.figure()
        for i in range(3):
            img=data[i][0].cpu().numpy().transpose([1, 2, 0])
            plt.subplot(1,3,i+1)
            plt.imshow(img)
            plt.axis('off')  # 不显示坐标轴
        plt.show()
        plt.close()


